#include "BatchProcessor.h"
#include <algorithm>

BatchProcessor::BatchProcessor(size_t maxBatchSize, std::chrono::milliseconds maxDelay)
    : m_maxBatchSize(maxBatchSize)
    , m_maxDelay(maxDelay)
    , m_shutdown(false)
    , m_totalOps(0)
    , m_batchedOps(0) {
    // 启动处理线程
    m_processorThread = std::thread(&BatchProcessor::ProcessorThread, this);
}

BatchProcessor::~BatchProcessor() {
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_shutdown = true;
    }
    m_condition.notify_one();
    
    if (m_processorThread.joinable()) {
        m_processorThread.join();
    }
}

void BatchProcessor::AddOperation(BatchOperation op) {
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_pendingOps.push(std::move(op));
        m_totalOps++;
    }
    m_condition.notify_one();
}

void BatchProcessor::ProcessorThread() {
    while (true) {
        std::vector<BatchOperation> batch;
        
        {
            std::unique_lock<std::mutex> lock(m_mutex);
            
            // 等待直到有操作需要处理或超时
            auto waitResult = m_condition.wait_for(lock, m_maxDelay, [this] {
                return m_shutdown || !m_pendingOps.empty();
            });
            
            if (m_shutdown && m_pendingOps.empty()) {
                return;
            }
            
            // 收集批处理操作
            while (!m_pendingOps.empty() && batch.size() < m_maxBatchSize) {
                batch.push_back(std::move(m_pendingOps.front()));
                m_pendingOps.pop();
            }
        }
        
        if (!batch.empty()) {
            ProcessBatchInternal(batch);
        }
    }
}

void BatchProcessor::ProcessBatchInternal(std::vector<BatchOperation>& batch) {
    // 按路径和偏移量排序
    std::sort(batch.begin(), batch.end());
    
    // 尝试合并操作
    if (TryMergeOperations(batch)) {
        m_batchedOps++;
    }
    
    // 处理每个批处理操作
    for (auto& op : batch) {
        HANDLE hFile = CreateFileW(
            op.path.c_str(),
            op.type == BatchOperation::READ ? GENERIC_READ : GENERIC_WRITE,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            nullptr,
            OPEN_EXISTING,
            FILE_FLAG_OVERLAPPED,
            nullptr
        );
        
        if (hFile == INVALID_HANDLE_VALUE) {
            op.callback(STATUS_UNSUCCESSFUL);
            continue;
        }
        
        OVERLAPPED overlapped = {0};
        overlapped.Offset = static_cast<DWORD>(op.offset & 0xFFFFFFFF);
        overlapped.OffsetHigh = static_cast<DWORD>(op.offset >> 32);
        
        BOOL success;
        if (op.type == BatchOperation::READ) {
            success = ReadFile(
                hFile,
                op.buffer.data(),
                static_cast<DWORD>(op.length),
                nullptr,
                &overlapped
            );
        } else {
            success = WriteFile(
                hFile,
                op.buffer.data(),
                static_cast<DWORD>(op.buffer.size()),
                nullptr,
                &overlapped
            );
        }
        
        if (!success && GetLastError() != ERROR_IO_PENDING) {
            CloseHandle(hFile);
            op.callback(STATUS_UNSUCCESSFUL);
            continue;
        }
        
        // 等待IO完成
        DWORD bytesTransferred;
        success = GetOverlappedResult(hFile, &overlapped, &bytesTransferred, TRUE);
        
        CloseHandle(hFile);
        op.callback(success ? STATUS_SUCCESS : STATUS_UNSUCCESSFUL);
    }
}

bool BatchProcessor::TryMergeOperations(std::vector<BatchOperation>& batch) {
    if (batch.size() < 2) return false;
    
    bool merged = false;
    auto it = batch.begin();
    while (it != batch.end() - 1) {
        auto next = it + 1;
        
        // 检查是否可以合并
        if (it->path == next->path &&
            it->type == next->type &&
            it->offset + it->length == next->offset) {
            
            // 合并缓冲区
            it->buffer.insert(it->buffer.end(),
                            next->buffer.begin(),
                            next->buffer.end());
            it->length += next->length;
            
            // 保存原始回调
            auto originalCallback = it->callback;
            auto nextCallback = next->callback;
            
            // 创建新的回调来处理两个操作
            it->callback = [originalCallback, nextCallback](NTSTATUS status) {
                originalCallback(status);
                nextCallback(status);
            };
            
            // 移除已合并的操作
            batch.erase(next);
            merged = true;
        } else {
            ++it;
        }
    }
    
    return merged;
}

void BatchProcessor::ProcessBatch() {
    std::vector<BatchOperation> batch;
    
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        while (!m_pendingOps.empty()) {
            batch.push_back(std::move(m_pendingOps.front()));
            m_pendingOps.pop();
        }
    }
    
    if (!batch.empty()) {
        ProcessBatchInternal(batch);
    }
}